package server

import (
	"encoding/hex"
	"errors"
	"fmt"
	"gitee.com/baidugoogle/rtsp-server/protocol/headers"
	"gitee.com/baidugoogle/rtsp-server/protocol/rtcpreceiver"
	"gitee.com/baidugoogle/rtsp-server/protocol/rtpaac"
	"gitee.com/baidugoogle/rtsp-server/protocol/rtph264"
	"github.com/general252/gout/ucolor"
	"github.com/pion/rtcp"
	"github.com/pion/rtp"
	"github.com/pion/sdp/v3"
	"github.com/pion/transport/packetio"
	"io"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// PusherTrack 一路推流
type PusherTrack struct {
	id     int32         // 唯一标识
	pusher *StreamPusher // 上级
	track  *Track        //

	isH264          bool
	h264Decoder     *rtph264.Decoder
	trackConfigH264 *TrackConfigH264
	fH264           *os.File // 测试

	isAAC      bool
	aacDecoder *rtpaac.Decoder

	ssrc         uint32                     // 客户端发送的 rtp ssrc
	rtcpReceiver *rtcpreceiver.RTCPReceiver // 接收者报告

	inTransport    *headers.Transport
	rtpClientPort  int // 客户端发送rtp的端口
	rtcpClientPort int // 客户端发送rtcp的端口

	rtcpAddr *net.UDPAddr

	rtpBuffer *packetio.Buffer // 接收缓冲区

	rtpBufferList    map[uint64]*packetio.Buffer // 发送缓冲区
	rtpBufferIndex   uint64
	rtpBufferListMux sync.RWMutex
}

func NewPusherTrack(track *Track, pusher *StreamPusher) *PusherTrack {
	c := &PusherTrack{
		id:     GetId(),
		pusher: pusher,
		track:  track,
		ssrc:   0,

		rtpBuffer:     packetio.NewBuffer(),
		rtpBufferList: map[uint64]*packetio.Buffer{},
	}

	c.rtpBuffer.SetLimitCount(25 * 60)

	if ssrcStr, ok := track.Media.Attribute(sdp.AttrKeySSRC); ok {
		vList := strings.Split(ssrcStr, " ")
		if len(vList) > 0 {
			if ssrcUint32, err := strconv.ParseUint(vList[0], 10, 32); err == nil {
				// SDP中ssrc
				c.SetSSRC(uint32(ssrcUint32))
			}
		}
	}

	clockRate, _ := c.track.ClockRate()
	c.rtcpReceiver = rtcpreceiver.New(time.Second*10, nil, clockRate, func(packet rtcp.Packet) {
		c.WriteRTCP(packet)
	})

	go c.run()

	c.isH264 = c.track.IsH264()
	c.isH264 = false
	if c.isH264 {
		c.trackConfigH264, _ = c.track.ExtractConfigH264()
		c.h264Decoder = rtph264.NewDecoder()
	}

	if c.isAAC = c.track.IsAAC(); c.isAAC {
		if aacCfg, err := c.track.ExtractConfigAAC(); err == nil {
			c.aacDecoder = rtpaac.NewDecoder(aacCfg.SampleRate)
		}
	}

	return c
}

func (c *PusherTrack) run() {

	packet := make([]byte, 1500)

	for {
		// 接收数据, 阻塞
		n, err := c.rtpBuffer.Read(packet)
		if err != nil {
			if err == io.EOF {
				return
			}

			log.Printf("read error %v", err)
			return
		}

		if n <= 0 {
			continue
		}

		// 内容分发
		var bufferList = map[uint64]*packetio.Buffer{}

		c.rtpBufferListMux.RLock()
		for id, buffer := range c.rtpBufferList {
			bufferList[id] = buffer
		}
		c.rtpBufferListMux.RUnlock()

		pkt := packet[:n]
		for id, buffer := range bufferList {
			if _, err = buffer.Write(pkt); err != nil {
				if err == io.ErrClosedPipe {
					c.RemoveRtpBuffer(id)
				} else {
					log.Println(err)
				}
			}
		}
	}
}

func (c *PusherTrack) Close() {
	_ = c.rtpBuffer.Close()
}

// Control streamid=[0/1]
func (c *PusherTrack) Control() string {
	v, ok := c.track.Media.Attribute("control")
	if !ok {
	}

	return v
}

func (c *PusherTrack) ClockRate() (int, error) {
	v, err := c.track.ClockRate()
	return v, err
}

func (c *PusherTrack) AddRtpBuffer(buf *packetio.Buffer) uint64 {
	c.rtpBufferListMux.Lock()
	defer c.rtpBufferListMux.Unlock()

	v := atomic.AddUint64(&c.rtpBufferIndex, 1)
	c.rtpBufferList[v] = buf

	return v
}

func (c *PusherTrack) RemoveRtpBuffer(id uint64) {
	c.rtpBufferListMux.Lock()
	defer c.rtpBufferListMux.Unlock()

	delete(c.rtpBufferList, id)
}

func (c *PusherTrack) WriteRTP(packet *rtp.Packet, rtpPacketData []byte) error {
	if true {
		c.rtcpReceiver.ProcessPacketRTP(time.Now(), packet, true) // 统计接收信息
		_, err := c.rtpBuffer.Write(rtpPacketData)
		return err
	}

	// pkt排序
	if c.isH264 {
		var pkt rtp.Packet
		if err := pkt.Unmarshal(rtpPacketData); err != nil {
			return err
		}

		// decode H264 NALUs from RTP packets
		nalus, pts, err := c.h264Decoder.DecodeUntilMarker(&pkt)
		if err != nil {
			if !errors.Is(err, rtph264.ErrMorePacketsNeeded) {
				log.Println(err)
			}
			return err
		}

		if c.fH264 == nil {
			c.fH264, err = os.OpenFile("E:/out.h264", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
		}

		for i := 0; i < len(nalus); i++ {
			nalu := nalus[i]

			if c.fH264 != nil {
				if NaluType(nalu[0]&0x1f) == nalu_type_idr {
					_, _ = c.fH264.Write([]byte{0, 0, 0, 1})
					_, _ = c.fH264.Write(c.trackConfigH264.SPS)
					_, _ = c.fH264.Write([]byte{0, 0, 0, 1})
					_, _ = c.fH264.Write(c.trackConfigH264.PPS)
				}

				_, _ = c.fH264.Write([]byte{0, 0, 0, 1})
				_, _ = c.fH264.Write(nalu)
			}

			if false {
				log.Printf("[h264] %v %v %-10v (%05.2fs)",
					hex.EncodeToString(nalu[:4]),
					getH264NalType(nalu[0]&0x1f),
					pts.Milliseconds(), pts.Seconds())
			}
		}
	}

	return nil
}

func (c *PusherTrack) OnRtcpPacket(pkt *RtcpPacket) {
	switch packet := pkt.Packet.(type) {
	case *rtcp.SenderReport:
		// 收到发送者报告
		c.rtcpReceiver.ProcessPacketRTCP(time.Now(), packet)
	}
}

func (c *PusherTrack) WriteRTCP(packet rtcp.Packet) {
	rtcpData, err := packet.Marshal()
	if err != nil {
		return
	}

	if c.inTransport.Protocol == headers.TransportProtocolTCP {
		_, _ = c.pusher.session.conn.Write(MarshalInterleavedFrame(c.inTransport.InterleavedIDs[0], rtcpData))
		return
	}

	if c.rtcpAddr == nil {
		addr, ok := c.pusher.session.GetClientAddress()
		if !ok {
			return
		}

		c.rtcpAddr = &net.UDPAddr{
			IP:   addr.IP,
			Port: c.rtcpClientPort,
		}
	}

	if c.rtcpAddr != nil {
		_, _ = c.GetRtcpSender().WriteToUDP(rtcpData, c.rtcpAddr)
	}
}

// SetSSRC 客户端发送的 rtp ssrc
func (c *PusherTrack) SetSSRC(ssrc uint32) {
	c.ssrc = ssrc
}

func (c *PusherTrack) GetRtpClientAddr() (string, bool) {
	addr, ok := c.pusher.session.GetClientAddress()
	if !ok {
		return "", false
	}

	return fmt.Sprintf("%v:%v", addr.IP.String(), c.rtpClientPort), true
}

func (c *PusherTrack) GetRtcpClientAddr() (string, bool) {
	addr, ok := c.pusher.session.GetClientAddress()
	if !ok {
		return "", false
	}

	return fmt.Sprintf("%v:%v", addr.IP.String(), c.rtcpClientPort), true
}

func (c *PusherTrack) GetRtpSender() *net.UDPConn {
	return c.pusher.session.server.GetRtpSender()
}

func (c *PusherTrack) GetRtcpSender() *net.UDPConn {
	return c.pusher.session.server.GetRtcpSender()
}

type NaluType byte

const (
	nalu_type_not_define NaluType = 0  // 没有定义,  // 1-23  NAL单元  单个 NAL 单元包 // 13-23 保留 // 30-31 没有定义
	nalu_type_slice      NaluType = 1  // 不分区，非IDR图像的片 slice_layer_without_partioning_rbsp() sliceheader
	nalu_type_dpa        NaluType = 2  // 片分区A slice_data_partition_a_layer_rbsp( ), slice_header
	nalu_type_dpb        NaluType = 3  // 片分区B slice_data_partition_b_layer_rbsp( )
	nalu_type_dpc        NaluType = 4  // 片分区C slice_data_partition_c_layer_rbsp( )
	nalu_type_idr        NaluType = 5  // IDR图像中的片 slice_layer_without_partitioning_rbsp( ),sliceheader
	nalu_type_sei        NaluType = 6  // 补充增强信息单元（SEI） sei_rbsp( )
	nalu_type_sps        NaluType = 7  // SPS seq_parameter_set_rbsp( )
	nalu_type_pps        NaluType = 8  // PPS pic_parameter_set_rbsp( )
	nalu_type_aud        NaluType = 9  // 序列结束 access_unit_delimiter_rbsp( )
	nalu_type_eoesq      NaluType = 10 // 序列结束 end_of_seq_rbsp( )
	nalu_type_eostream   NaluType = 11 // 码流借宿 end_of_stream_rbsp( )
	nalu_type_filler     NaluType = 12 // 填充 filler_data_rbsp( )

	nalu_type_stap_a NaluType = 24 // 单一时间的组合包
	nalu_type_stap_b NaluType = 25 // 单一时间的组合包
	nalu_type_MTAP16 NaluType = 26 // 多个时间的组合包
	nalu_type_MTAP24 NaluType = 27 // 多个时间的组合包
	nalu_type_fu_a   NaluType = 28 // FU-A 分片的单元
	nalu_type_fu_b   NaluType = 29 // FU-B 分片的单元
)

var getH264NalType = func(n byte) string {
	nType := n & 0x1f
	switch NaluType(nType) {
	case nalu_type_slice:
		return ucolor.Yellow("slice")
	case nalu_type_aud:
		return "aud"
	case nalu_type_idr:
		return ucolor.HiGreen("IDR")
	case nalu_type_sei:
		return ucolor.HiRed("SEI")
	case nalu_type_sps:
		return "SPS"
	case nalu_type_pps:
		return "PPS"
	case nalu_type_fu_a:
		return "FU-A"
	default:
		return fmt.Sprintf("%v", nType)
	}
}
